use async_stream;
use futures::stream::{Stream, StreamExt};
use rig::{completion::ToolDefinition, tool::Tool};
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::RwLock;

use crate::agent_types::AgentCategory;
use crate::agents::{AgentManager, ChatMessage};
use crate::embedding::EmbeddingManager;
use crate::image::ImageManager;
use crate::tools::ToolManager;

/// Agent orchestration modes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum OrchestrationMode {
    /// Sequential execution of agents
    Sequential,
    /// Parallel execution of agents
    Parallel,
    /// Adaptive routing based on input
    Adaptive,
    /// Multi-turn conversation
    MultiTurn,
}

/// Agent workflow configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowConfig {
    /// Workflow name
    pub name: String,
    /// Description of the workflow
    pub description: String,
    /// Orchestration mode
    pub mode: OrchestrationMode,
    /// Agent IDs involved in the workflow
    pub agent_ids: Vec<i64>,
    /// Workflow steps (for sequential mode)
    pub steps: Option<Vec<WorkflowStep>>,
    /// Routing rules (for adaptive mode)
    pub routing_rules: Option<Vec<RoutingRule>>,
}

/// Individual workflow step
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowStep {
    /// Step name
    pub name: String,
    /// Agent ID for this step
    pub agent_id: i64,
    /// Step prompt template
    pub prompt_template: String,
    /// Whether to pass previous step output
    pub pass_output: bool,
    /// Maximum retries
    pub max_retries: u32,
}

/// Routing rule for adaptive workflows
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingRule {
    /// Rule name
    pub name: String,
    /// Condition to match
    pub condition: String,
    /// Agent ID to route to
    pub target_agent_id: i64,
    /// Confidence threshold
    pub confidence_threshold: f32,
}

/// Workflow execution result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowResult {
    /// Workflow name
    pub workflow_name: String,
    /// Execution results per step
    pub step_results: Vec<StepResult>,
    /// Final output
    pub final_output: String,
    /// Execution metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

/// Individual step result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StepResult {
    /// Step name
    pub step_name: String,
    /// Agent ID used
    pub agent_id: i64,
    /// Input prompt
    pub input: String,
    /// Output result
    pub output: String,
    /// Execution time in milliseconds
    pub execution_time_ms: u64,
    /// Whether step succeeded
    pub success: bool,
    /// Error message if failed
    pub error_message: Option<String>,
}

/// Advanced agent manager for complex workflows
#[derive(Debug)]
pub struct AdvancedAgentManager {
    /// Base agent manager
    agent_manager: Arc<AgentManager>,
    /// Tool manager
    tool_manager: Arc<ToolManager>,
    /// Embedding manager
    embedding_manager: Arc<EmbeddingManager>,
    /// Image manager
    image_manager: Arc<ImageManager>,
    /// Workflow configurations
    workflows: Arc<RwLock<HashMap<String, WorkflowConfig>>>,
}

impl AdvancedAgentManager {
    /// Create a new advanced agent manager
    pub fn new(
        agent_manager: Arc<AgentManager>,
        tool_manager: Arc<ToolManager>,
        embedding_manager: Arc<EmbeddingManager>,
        image_manager: Arc<ImageManager>,
    ) -> Self {
        Self {
            agent_manager,
            tool_manager,
            embedding_manager,
            image_manager,
            workflows: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a new workflow
    pub async fn register_workflow(&self, workflow: WorkflowConfig) -> anyhow::Result<()> {
        let mut workflows = self.workflows.write().await;
        workflows.insert(workflow.name.clone(), workflow);
        Ok(())
    }

    /// Execute a workflow
    pub async fn execute_workflow(
        &self,
        workflow_name: &str,
        input: &str,
    ) -> anyhow::Result<WorkflowResult> {
        let workflows = self.workflows.read().await;
        let workflow = workflows
            .get(workflow_name)
            .ok_or_else(|| anyhow::anyhow!("Workflow '{}' not found", workflow_name))?
            .clone();
        drop(workflows);

        match workflow.mode {
            OrchestrationMode::Sequential => {
                self.execute_sequential_workflow(&workflow, input).await
            }
            OrchestrationMode::Parallel => self.execute_parallel_workflow(&workflow, input).await,
            OrchestrationMode::Adaptive => self.execute_adaptive_workflow(&workflow, input).await,
            OrchestrationMode::MultiTurn => {
                self.execute_multi_turn_workflow(&workflow, input).await
            }
        }
    }

    /// Execute sequential workflow
    async fn execute_sequential_workflow(
        &self,
        workflow: &WorkflowConfig,
        input: &str,
    ) -> anyhow::Result<WorkflowResult> {
        let mut step_results = Vec::new();
        let mut current_output = input.to_string();

        if let Some(steps) = &workflow.steps {
            for step in steps {
                let start_time = std::time::Instant::now();

                // Prepare prompt for this step
                let prompt = if step.pass_output && !step_results.is_empty() {
                    step.prompt_template.replace("{input}", &current_output)
                } else {
                    step.prompt_template.replace("{input}", input)
                };

                // Execute step with retry logic
                let mut result = String::new();
                let mut success = false;
                let mut error_message = None;

                for attempt in 0..=step.max_retries {
                    match crate::completion::prompt(
                        &self.agent_manager,
                        &self.tool_manager,
                        step.agent_id,
                        &prompt,
                    )
                    .await
                    {
                        Ok(output) => {
                            result = output;
                            success = true;
                            break;
                        }
                        Err(e) => {
                            if attempt == step.max_retries {
                                error_message = Some(e.to_string());
                            }
                            // Wait before retry
                            tokio::time::sleep(tokio::time::Duration::from_millis(
                                100 * (attempt + 1) as u64,
                            ))
                            .await;
                        }
                    }
                }

                let execution_time_ms = start_time.elapsed().as_millis() as u64;

                step_results.push(StepResult {
                    step_name: step.name.clone(),
                    agent_id: step.agent_id,
                    input: prompt.clone(),
                    output: result.clone(),
                    execution_time_ms,
                    success,
                    error_message,
                });

                current_output = result;
            }
        }

        Ok(WorkflowResult {
            workflow_name: workflow.name.clone(),
            step_results,
            final_output: current_output,
            metadata: HashMap::new(),
        })
    }

    /// Execute parallel workflow
    async fn execute_parallel_workflow(
        &self,
        workflow: &WorkflowConfig,
        input: &str,
    ) -> anyhow::Result<WorkflowResult> {
        let mut tasks = Vec::new();

        for &agent_id in &workflow.agent_ids {
            let agent_manager = self.agent_manager.clone();
            let tool_manager = self.tool_manager.clone();
            let input = input.to_string();

            tasks.push(tokio::spawn(async move {
                let start_time = std::time::Instant::now();
                let result =
                    crate::completion::prompt(&agent_manager, &tool_manager, agent_id, &input)
                        .await;

                let execution_time_ms = start_time.elapsed().as_millis() as u64;

                StepResult {
                    step_name: format!("Agent_{}", agent_id),
                    agent_id,
                    input: input.clone(),
                    output: result
                        .as_ref()
                        .map(|s| s.clone())
                        .unwrap_or_else(|_| "error".to_string()),
                    execution_time_ms,
                    success: result.is_ok(),
                    error_message: result.as_ref().err().map(|e| e.to_string()),
                }
            }));
        }

        let step_results: Vec<StepResult> = futures::future::join_all(tasks)
            .await
            .into_iter()
            .map(|r| {
                r.unwrap_or_else(|e| StepResult {
                    step_name: "Error".to_string(),
                    agent_id: 0,
                    input: String::new(),
                    output: String::new(),
                    execution_time_ms: 0,
                    success: false,
                    error_message: Some(e.to_string()),
                })
            })
            .collect();

        // Combine all outputs for final result
        let final_output = step_results
            .iter()
            .filter(|r| r.success)
            .map(|r| format!("Agent {}: {}\n", r.agent_id, r.output))
            .collect::<String>();

        Ok(WorkflowResult {
            workflow_name: workflow.name.clone(),
            step_results,
            final_output,
            metadata: HashMap::new(),
        })
    }

    /// Execute adaptive workflow
    async fn execute_adaptive_workflow(
        &self,
        workflow: &WorkflowConfig,
        input: &str,
    ) -> anyhow::Result<WorkflowResult> {
        let mut step_results = Vec::new();
        let mut current_output = input.to_string();
        let mut current_agent_id = workflow
            .agent_ids
            .first()
            .copied()
            .ok_or_else(|| anyhow::anyhow!("No agents specified for workflow"))?;

        // Simple adaptive routing (in real implementation, you'd use more sophisticated routing)
        for iteration in 0..5 {
            let start_time = std::time::Instant::now();

            let result = crate::completion::prompt(
                &self.agent_manager,
                &self.tool_manager,
                current_agent_id,
                &current_output,
            )
            .await;

            let execution_time_ms = start_time.elapsed().as_millis() as u64;

            step_results.push(StepResult {
                step_name: format!("Adaptive_Step_{}", iteration + 1),
                agent_id: current_agent_id,
                input: current_output.clone(),
                output: result
                    .as_ref()
                    .map(|s| s.clone())
                    .unwrap_or_else(|_| "error".to_string()),
                execution_time_ms,
                success: result.is_ok(),
                error_message: result.as_ref().err().map(|e| e.to_string()),
            });

            if let Ok(output) = result {
                current_output = output;

                // Simple routing logic - in real implementation, use the routing rules
                if current_output.to_lowercase().contains("continue") {
                    // Move to next agent
                    if let Some(index) = workflow
                        .agent_ids
                        .iter()
                        .position(|&id| id == current_agent_id)
                    {
                        if let Some(&next_id) = workflow.agent_ids.get(index + 1) {
                            current_agent_id = next_id;
                        }
                    }
                } else {
                    break;
                }
            } else {
                break;
            }
        }

        Ok(WorkflowResult {
            workflow_name: workflow.name.clone(),
            step_results,
            final_output: current_output,
            metadata: HashMap::new(),
        })
    }

    /// Execute multi-turn workflow
    async fn execute_multi_turn_workflow(
        &self,
        workflow: &WorkflowConfig,
        input: &str,
    ) -> anyhow::Result<WorkflowResult> {
        let mut step_results = Vec::new();
        let mut history = Vec::new();
        let mut current_input = input.to_string();

        for turn in 0..5 {
            let start_time = std::time::Instant::now();

            // Use the first agent for multi-turn conversation
            let agent_id = workflow
                .agent_ids
                .first()
                .copied()
                .ok_or_else(|| anyhow::anyhow!("No agents specified for workflow"))?;

            let result = crate::completion::chat(
                &self.agent_manager,
                &self.tool_manager,
                agent_id,
                &current_input,
                history.clone(),
            )
            .await;

            let execution_time_ms = start_time.elapsed().as_millis() as u64;

            step_results.push(StepResult {
                step_name: format!("Turn_{}", turn + 1),
                agent_id,
                input: current_input.clone(),
                output: result
                    .as_ref()
                    .map(|s| s.clone())
                    .unwrap_or_else(|_| "error".to_string()),
                execution_time_ms,
                success: result.is_ok(),
                error_message: result.as_ref().err().map(|e| e.to_string()),
            });

            if let Ok(response) = result {
                history.push(ChatMessage {
                    role: "user".to_string(),
                    content: current_input,
                });
                history.push(ChatMessage {
                    role: "assistant".to_string(),
                    content: response.clone(),
                });

                current_input =
                    "Please continue the conversation or say 'done' to end.".to_string();

                if response.to_lowercase().contains("done") {
                    break;
                }
            } else {
                break;
            }
        }

        let final_output = history
            .iter()
            .filter(|msg| msg.role == "assistant")
            .map(|msg| msg.content.clone())
            .collect::<Vec<_>>()
            .join("\n");

        Ok(WorkflowResult {
            workflow_name: workflow.name.clone(),
            step_results,
            final_output,
            metadata: HashMap::new(),
        })
    }

    /// Create a specialized agent with tools
    pub async fn create_specialized_agent(
        &self,
        name: &str,
        provider: &str,
        model: &str,
        preamble: &str,
        category: AgentCategory,
        tools: Option<Vec<String>>,
    ) -> anyhow::Result<i64> {
        self.agent_manager
            .create_agent(
                name,
                provider,
                model,
                preamble,
                None,
                tools.as_deref(),
                category,
                None,
            )
            .await
    }

    /// Execute RAG query with context retrieval
    pub async fn execute_rag_query(
        &self,
        agent_id: i64,
        query: &str,
        context_docs: usize,
    ) -> anyhow::Result<String> {
        crate::embedding::rag_query(
            &self.agent_manager,
            &self.embedding_manager,
            agent_id,
            query,
            context_docs,
        )
        .await
    }

    /// Stream workflow execution
    pub async fn stream_workflow_execution<'a>(
        &'a self,
        workflow_name: &str,
        input: &'a str,
    ) -> anyhow::Result<Pin<Box<dyn Stream<Item = anyhow::Result<String>> + Send + 'a>>> {
        let workflows = self.workflows.read().await;
        let workflow = workflows
            .get(workflow_name)
            .ok_or_else(|| anyhow::anyhow!("Workflow '{}' not found", workflow_name))?
            .clone();
        drop(workflows);

        // Create a stream that emits workflow execution progress
        let stream = async_stream::stream! {
            yield Ok(format!("Starting workflow: {}", workflow.name));

            if let Some(steps) = &workflow.steps {
                for (index, step) in steps.iter().enumerate() {
                    yield Ok(format!("Executing step {}: {}", index + 1, step.name));

                    let prompt = step.prompt_template.replace("{input}", input);

                    match crate::completion::prompt(
                        &self.agent_manager,
                        &self.tool_manager,
                        step.agent_id,
                        &prompt,
                    )
                    .await
                    {
                        Ok(output) => {
                            yield Ok(format!("Step {} completed: {}", index + 1, step.name));
                            yield Ok(format!("Output: {}", output));
                        }
                        Err(e) => {
                            yield Ok(format!("Step {} failed: {}", index + 1, e));
                        }
                    }
                }
            }

            yield Ok(format!("Workflow {} completed", workflow.name));
        };

        Ok(Box::pin(stream))
    }
}

/// Built-in tools for advanced agents
#[derive(Deserialize, Serialize)]
pub struct WorkflowToolArgs {
    pub workflow_name: String,
    pub input: String,
}

#[derive(Debug, thiserror::Error)]
#[error("Workflow tool error")]
pub struct WorkflowToolError;

#[derive(Deserialize, Serialize)]
pub struct WorkflowTool;

impl Tool for WorkflowTool {
    const NAME: &'static str = "execute_workflow";
    type Error = WorkflowToolError;
    type Args = WorkflowToolArgs;
    type Output = String;

    async fn definition(&self, _prompt: String) -> ToolDefinition {
        ToolDefinition {
            name: "execute_workflow".to_string(),
            description: "Execute a registered workflow with the given input".to_string(),
            parameters: json!({
                "type": "object",
                "properties": {
                    "workflow_name": {
                        "type": "string",
                        "description": "Name of the workflow to execute"
                    },
                    "input": {
                        "type": "string",
                        "description": "Input for the workflow"
                    }
                },
                "required": ["workflow_name", "input"]
            }),
        }
    }

    async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
        // This would need to be implemented with access to the AdvancedAgentManager
        // For now, return a placeholder response
        Ok(format!(
            "Workflow '{}' executed with input: '{}'",
            args.workflow_name, args.input
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[tokio::test]
    async fn test_advanced_agent_manager_creation() {
        let temp_dir = TempDir::new().unwrap();
        let agents_file = temp_dir.path().join("agents.json");
        let mcp_file = temp_dir.path().join("mcp.json");
        let agent_manager = Arc::new(AgentManager::new(&agents_file).unwrap());
        let tool_manager = Arc::new(ToolManager::new(&mcp_file));
        let embedding_manager = Arc::new(EmbeddingManager::new());
        let image_manager = Arc::new(ImageManager::new());

        let advanced_manager = AdvancedAgentManager::new(
            agent_manager,
            tool_manager,
            embedding_manager,
            image_manager,
        );

        // Test that the manager was created successfully
        assert!(true);
    }

    #[tokio::test]
    async fn test_workflow_registration() {
        let temp_dir = TempDir::new().unwrap();
        let agents_file = temp_dir.path().join("agents.json");
        let agent_manager = Arc::new(AgentManager::new(&agents_file).unwrap());
        let tool_manager = Arc::new(ToolManager::new());
        let embedding_manager = Arc::new(EmbeddingManager::new());
        let image_manager = Arc::new(ImageManager::new());

        let advanced_manager = AdvancedAgentManager::new(
            agent_manager,
            tool_manager,
            embedding_manager,
            image_manager,
        );

        let workflow = WorkflowConfig {
            name: "test_workflow".to_string(),
            description: "Test workflow".to_string(),
            mode: OrchestrationMode::Sequential,
            agent_ids: vec![1],
            steps: None,
            routing_rules: None,
        };

        let result = advanced_manager.register_workflow(workflow).await;
        assert!(result.is_ok());
    }
}
